from pytorch_metric_learning import losses, miners
from pytorch_metric_learning.distances import CosineSimilarity, DotProductSimilarity
from .TripletCenterLoss import TripletCenterLoss

def get_loss_miner_cfg(cfg):
    loss = get_loss_cfg(cfg)
    miner_name = cfg.MODEL.MASK_FORMER.INS_MINER
    miner = get_miner(miner_name)
    return loss,miner

def get_loss_cfg(cfg):
    loss_name = cfg.MODEL.MASK_FORMER.INS_LOSS
    if loss_name == "TripletCenterLoss":
        mg = cfg.MODEL.MASK_FORMER.INS_LOSS_MARGIN
        dim = cfg.MODEL.MASK_FORMER.INS_LOSS_DIM
        nc = cfg.MODEL.MASK_FORMER.INS_LOSS_NUM
        ins_loss = get_triplet_center_loss(margin=mg,numclass=nc,dim=dim)
        return ins_loss
    if loss_name == "ArcFaceLoss":
        margin=cfg.MODEL.MASK_FORMER.INS_LOSS_MARGIN
        numclass = cfg.MODEL.MASK_FORMER.INS_LOSS_NUM
        dim=cfg.MODEL.MASK_FORMER.INS_LOSS_DIM
        return losses.ArcFaceLoss(numclass,embedding_size=dim,margin=margin)
    if loss_name == 'SupConLoss': return losses.SupConLoss(temperature=0.07)
    if loss_name == 'CircleLoss': return losses.CircleLoss(m=0.4, gamma=80) #these are params for image retrieval
    if loss_name == 'MultiSimilarityLoss': return losses.MultiSimilarityLoss(alpha=1.0, beta=50, base=0.0, distance=DotProductSimilarity())
    if loss_name == 'ContrastiveLoss': return losses.ContrastiveLoss(pos_margin=0, neg_margin=1)
    if loss_name == 'Lifted': return losses.GeneralizedLiftedStructureLoss(neg_margin=0, pos_margin=1, distance=DotProductSimilarity())
    if loss_name == 'FastAPLoss': return losses.FastAPLoss(num_bins=30)
    if loss_name == 'NTXentLoss': return losses.NTXentLoss(temperature=0.07) #The MoCo paper uses 0.07, while SimCLR uses 0.5.
    if loss_name == 'TripletMarginLoss': return losses.TripletMarginLoss(margin=0.1, swap=False, smooth_loss=False, triplets_per_anchor='all') #or an int, for example 100
    if loss_name == 'ArcFaceLoss': return losses.ArcFaceLoss()
    if loss_name == 'CentroidTripletLoss': return losses.CentroidTripletLoss(margin=0.05,
                                                                            swap=False,
                                                                            smooth_loss=False,
                                                                            triplets_per_anchor="all",)
    return None

def get_triplet_center_loss(margin=0.5,numclass = 200,dim=4096):
    return TripletCenterLoss(margin=margin,num_classes=numclass,dim=dim)

def get_arcface_loss(opt):
    margin=opt.margin
    numclass = 200
    dim=4096
    return losses.ArcFaceLoss(numclass,embedding_size=dim,margin=margin)

def get_loss(loss_name):
    # print("INS Loss is %s" % (loss_name))
    if loss_name == 'SupConLoss': return losses.SupConLoss(temperature=0.07)
    if loss_name == 'CircleLoss': return losses.CircleLoss(m=0.4, gamma=80) #these are params for image retrieval
    if loss_name == 'MultiSimilarityLoss': return losses.MultiSimilarityLoss(alpha=1.0, beta=50, base=0.0, distance=DotProductSimilarity())
    if loss_name == 'ContrastiveLoss': return losses.ContrastiveLoss(pos_margin=0, neg_margin=1)
    if loss_name == 'Lifted': return losses.GeneralizedLiftedStructureLoss(neg_margin=0, pos_margin=1, distance=DotProductSimilarity())
    if loss_name == 'FastAPLoss': return losses.FastAPLoss(num_bins=30)
    if loss_name == 'NTXentLoss': return losses.NTXentLoss(temperature=0.07) #The MoCo paper uses 0.07, while SimCLR uses 0.5.
    if loss_name == 'TripletMarginLoss': return losses.TripletMarginLoss(margin=0.1, swap=False, smooth_loss=False, triplets_per_anchor='all') #or an int, for example 100
    if loss_name == 'ArcFaceLoss': return losses.ArcFaceLoss()
    if loss_name == 'CentroidTripletLoss': return losses.CentroidTripletLoss(margin=0.05,
                                                                            swap=False,
                                                                            smooth_loss=False,
                                                                            triplets_per_anchor="all",)
    return loss_name

def get_miner(miner_name, margin=0.1):
    print("INS MINER is %s" % (miner_name))
    if miner_name == 'TripletMarginMiner' : return miners.TripletMarginMiner(margin=margin, type_of_triplets="semihard") # all, hard, semihard, easy    
    if miner_name == 'MultiSimilarityMiner' : return miners.MultiSimilarityMiner(epsilon=margin, distance=CosineSimilarity())
    if miner_name == 'PairMarginMiner' : return miners.PairMarginMiner(pos_margin=0.7, neg_margin=0.3, distance=DotProductSimilarity())
    return None
